'''
ManiSkill 环境测试 - 修复版
'''

import gymnasium as gym
import mani_skill.envs
import numpy as np
import torch

# 列出可用的 ManiSkill 环境
print("Available ManiSkill environments:")
available_envs = [env_id for env_id in gym.envs.registry.keys() if any(keyword in env_id for keyword in ['MS', 'Cube', 'Pick', 'Push', 'Stack'])]
for env_id in available_envs[:15]:  # 显示前15个
    print(f"  {env_id}")

def safe_format_reward(reward):
    """安全地格式化奖励值，处理 tensor 类型"""
    if isinstance(reward, torch.Tensor):
        return float(reward.item())
    elif isinstance(reward, np.ndarray):
        return float(reward.item()) if reward.size == 1 else float(reward.sum())
    else:
        return float(reward)

def test_environment(env_id, episodes=3):
    """测试指定的 ManiSkill 环境"""
    try:
        print(f"\n=== Testing {env_id} ===")
        
        # 创建环境
        env = gym.make(env_id, render_mode="rgb_array")
        print(f"Environment created successfully!")
        print(f"Observation space: {env.observation_space}")
        print(f"Action space: {env.action_space}")
        
        max_steps = getattr(env.spec, 'max_episode_steps', None)
        if max_steps is None:
            max_steps = 200  # 设置默认最大步数
        print(f"Max episode steps: {max_steps}")
        
        total_reward = 0
        total_success = 0
        
        for episode in range(episodes):
            print(f"\n--- Episode {episode + 1} ---")
            
            try:
                obs, info = env.reset()
                episode_reward = 0
                episode_steps = 0
                
                for step in range(max_steps):
                    # 随机动作
                    action = env.action_space.sample()
                    
                    # 执行动作
                    obs, reward, terminated, truncated, info = env.step(action)
                    
                    # 安全地处理奖励值
                    reward_val = safe_format_reward(reward)
                    episode_reward += reward_val
                    episode_steps += 1
                    
                    # 打印有用信息 (只在奖励非零时)
                    if abs(reward_val) > 1e-6:
                        print(f"  Step {episode_steps}: reward = {reward_val:.3f}")
                    
                    # 检查是否成功完成任务
                    if isinstance(info, dict) and 'success' in info and info['success']:
                        print(f"  Task completed successfully at step {episode_steps}!")
                        total_success += 1
                        break
                    
                    # 检查是否结束
                    if terminated or truncated:
                        break
                
                print(f"  Episode ended: steps = {episode_steps}, total_reward = {episode_reward:.3f}")
                total_reward += episode_reward
                
            except Exception as e:
                print(f"  Error in episode {episode + 1}: {e}")
                continue
        
        # 统计结果
        mean_reward = total_reward / episodes if episodes > 0 else 0
        success_rate = (total_success / episodes * 100) if episodes > 0 else 0
        
        print(f"\n--- Results for {env_id} ---")
        print(f"Episodes completed: {episodes}")
        print(f"Mean reward: {mean_reward:.3f}")
        print(f"Success rate: {success_rate:.1f}%")
        
        env.close()
        return True
        
    except Exception as e:
        print(f"Failed to test {env_id}: {e}")
        return False

# 测试环境列表 - 使用实际可用的环境
test_envs = [
    "PushCube-v1",     # 这个在上面的输出中成功了
    "SceneManipulation-v1",  # 从可用环境列表中选择
]

# 如果找不到 Cube 相关环境，添加更多备选
additional_envs = []
for env_id in gym.envs.registry.keys():
    if any(keyword in env_id.lower() for keyword in ['pick', 'push', 'stack', 'lift']) and 'v1' in env_id:
        additional_envs.append(env_id)
        if len(additional_envs) >= 3:  # 限制数量
            break

test_envs.extend(additional_envs)

print(f"\nTesting {len(set(test_envs))} environments...")
successful_envs = []

for env_id in list(set(test_envs)):  # 去重
    if test_environment(env_id, episodes=2):  # 减少到2个episode
        successful_envs.append(env_id)

print(f"\n=== Final Summary ===")
print(f"Successfully tested environments: {successful_envs}")
print(f"Total successful: {len(successful_envs)}")

# 如果有成功的环境，展示一个详细示例
if successful_envs:
    print(f"\n=== Detailed Example with {successful_envs[0]} ===")
    try:
        env = gym.make(successful_envs[0], render_mode="rgb_array")
        obs, info = env.reset()
        
        print(f"Initial observation shape: {obs.shape if hasattr(obs, 'shape') else type(obs)}")
        print(f"Initial info: {info}")
        
        # 执行几步
        for i in range(5):
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            reward_val = safe_format_reward(reward)
            print(f"Step {i+1}: reward = {reward_val:.4f}, terminated = {terminated}, truncated = {truncated}")
            
            if terminated or truncated:
                break
        
        env.close()
        
    except Exception as e:
        print(f"Error in detailed example: {e}")